import torch
import torchvision as torchvision
from torch import nn
from torchvision.transforms import ToPILImage
from torch.autograd import Variable # 图像处理有关函数
from net import MyLeNet5  # 导入搭建的网络
from torch.optim import lr_scheduler  # 学习率 优化器
from torchvision import datasets, transforms  # 数据集加载
from torch.utils.data import DataLoader, Dataset
import os

# 数据转换为tensor格式
data_transform = transforms.Compose([
    transforms.ToTensor()
])

# 加载内置的训练数据集：minst
train_dataset = datasets.MNIST(root="./data", train=True, transform=data_transform, download=True)
train_dataloader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=16, shuffle=True)
# 加载内置的训练数据集：minst
test_dataset = datasets.MNIST(root="./data", train=False, transform=data_transform, download=True)
test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=16, shuffle=True)

# 如果有显卡，可以转到GPU
device = "cuda" if torch.cuda.is_available() else 'cpu'

# 调用之前搭建好的网络模型，将模型数据转到GPU
model = MyLeNet5().to(device)

model.load_state_dict(torch.load("D:/PyCharm Coder Files/LeNet-5/save_model/best_model.pth"))   # 加载权重文件

# 获取结果
classes = [
    "0",
    "1",
    "2",
    "3",
    "4",
    "5",
    "6",
    "7",
    "8",
    "9"
]

# 把tensor转换为图片，方便可视化
show = ToPILImage()

# 进行测试
s = 0
epoch = 200  # 测试次数

for i in range(epoch):
    X, y = test_dataset[i][0], test_dataset[i][1]   # X为图片，y为标签
    # show(X).show()    # 展示图片
    # 将张量拓展为4维 X原来为3维：CHW
    # 对数据的维度进行压缩squeeze或者解压unsqueeze ,dim = 0表示行 ，dim = 1 表示列
    X = Variable(torch.unsqueeze(X, dim=0).float(), requires_grad=False).to(device)
    # X.shape==torch.Size([3]) ===>> X.shape==torch.Size([1, 3])
    with torch.no_grad():
        pred = model(X)  # model要求四维，BCHW
        predicted, actual = classes[torch.argmax(pred[0])], classes[y]

        s += 1 if classes[torch.argmax(pred[0])] == (classes[y]) else 0
        print(f'predicted:{predicted},actual:{actual}')

acc = s / epoch
print(f"accuracy:{acc}")

